#include "vm.h"
#include "pmm.h"
#include "../../include/memlayout.h"
#include "../../include/riscv.h"
#include "../boot/printf.h"
#include "../proc/proc.h"
#include "../file/string.h"
#define MEM_START   0x80000000UL
#define MEM_END     0x88000000UL
extern char trampoline[];
// 内核页表
pagetable_t kernel_pagetable = 0;
extern char etext[];
static inline int is_page_aligned(uint64 addr) {
    return (addr & (PGSIZE - 1)) == 0;
}

static pte_t* walk_inner(pagetable_t pt, uint64 va, int alloc) {
    // 仅支持 39-bit VA
    if (va >> 39) return 0;
    pagetable_t pagetable = pt;
    for (int level = 2; level > 0; level--) {
        uint64 idx = VPN_MASK(va, level);
        pte_t *pte = &pagetable[idx];
        if (*pte & PTE_V) {
            // 前进到下一级页表
            pagetable = (pagetable_t)PTE2PA(*pte);
        } else {
            if (!alloc) return 0;
            void *newpage = alloc_page();
            if (!newpage) return 0;
            // 清零新页表
            for (int i = 0; i < 512; ++i) ((uint64*)newpage)[i] = 0;
            *pte = PA2PTE(newpage) | PTE_V;
            pagetable = (pagetable_t)newpage;
        }
    }
    // 到达最低级
    return &pagetable[VPN_MASK(va, 0)];
}

pte_t* walk_create(pagetable_t pt, uint64 va) { return walk_inner(pt, va, 1); }
pte_t* walk_lookup(pagetable_t pt, uint64 va) { return walk_inner(pt, va, 0); }

pagetable_t create_pagetable(void) {
    void *p = alloc_page();
    if (!p) return 0;
    for (int i = 0; i < 512; ++i) ((uint64*)p)[i] = 0;
    return (pagetable_t)p;
}

static void free_walk(pagetable_t pt, int level) {
    if (!pt) return;
    if (level == 0) return; // L0 由上层根据映射拆除后释放整页表
    for (int i = 0; i < 512; ++i) {
        pte_t pte = pt[i];
        if ((pte & PTE_V) && !(pte & (PTE_R|PTE_W|PTE_X))) {
            // 指向下一级页表
            pagetable_t child = (pagetable_t)PTE2PA(pte);
            free_walk(child, level - 1);
            pt[i] = 0;
            free_page(child);
        }
    }
}

void destroy_pagetable(pagetable_t pt) {
    if (!pt) return;
    free_walk(pt, 2);
    free_page(pt);
}

int map_page(pagetable_t pt, uint64 va, uint64 pa, int perm) {
    if (!is_page_aligned(va) || !is_page_aligned(pa)) return -1;
    pte_t *pte = walk_create(pt, va);
    if (!pte) return -1;
    if (*pte & PTE_V) return -2; // 冲突
    *pte = PA2PTE(pa) | perm | PTE_V | PTE_A | ((perm & PTE_W) ? PTE_D : 0);
    return 0;
}

uint64
uvmalloc(pagetable_t pagetable, uint64 oldsz, uint64 newsz, int perm)
{
    if(newsz < oldsz)
        return oldsz;

    uint64 start = PGROUNDUP(oldsz);
    uint64 end = PGROUNDUP(newsz);
    int flags = perm | PTE_R | PTE_U;
    uint64 a;

    for(a = start; a < end; a += PGSIZE){
        void *mem = alloc_page();
        if(mem == 0)
            goto err;
        memset(mem, 0, PGSIZE);
        if(map_page(pagetable, a, (uint64)mem, flags) != 0){
            free_page(mem);
            goto err;
        }
    }

    return newsz;

err:
    if(a > start)
        uvmunmap(pagetable, start, (a - start) / PGSIZE, 1);
    return 0;
}

static void kvmmap_range(pagetable_t pt, uint64 va, uint64 pa, uint64 sz, int perm) {
    uint64 a = PGROUNDDOWN(va);
    uint64 last = PGROUNDDOWN(va + sz - 1);
    for (;; a += PGSIZE, pa += PGSIZE) {
        map_page(pt, a, pa, perm);
        if (a == last) break;
    }
}
#define PLIC_BASE 0x0c000000ul
#define CLINT_BASE 0x2000000ul
void kvminit(void) {
    kernel_pagetable = create_pagetable();
    kvmmap_range(kernel_pagetable, PLIC_BASE, PLIC_BASE, PGSIZE, PTE_R|PTE_W);
    kvmmap_range(kernel_pagetable, CLINT_BASE, CLINT_BASE, PGSIZE, PTE_R|PTE_W);
    kvmmap_range(kernel_pagetable, UART0, UART0, PGSIZE, PTE_R|PTE_W);
    kvmmap_range(kernel_pagetable, MEM_START, MEM_START, MEM_END - MEM_START, PTE_R | PTE_W|PTE_X);

    // 映射trampoline到内核高地址空间，供forkret/usertrap返回路径使用
    if(map_page(kernel_pagetable, TRAMPOLINE, (uint64)trampoline, PTE_R | PTE_X)!=0){
        panic("Failed to map trampoline");
    };
    
    // 映射所有进程的内核栈
    proc_mapstacks(kernel_pagetable);
}

void kvminithart(void) {
    uint64 satp = MAKE_SATP(kernel_pagetable);
    // 设置satp寄存器
    w_satp(satp);
    sfence_vma();
}

// 取消映射用户地址空间中的指定页
// va: 起始虚拟地址（必须是页对齐的）
// npages: 要取消映射的页数
// do_free: 如果为1，释放对应的物理内存；如果为0，只取消映射
void uvmunmap(pagetable_t pagetable, uint64 va, uint64 npages, int do_free) {
    uint64 a;
    pte_t *pte;

    if ((va % PGSIZE) != 0)
        panic("uvmunmap: not aligned");

    for (a = va; a < va + npages * PGSIZE; a += PGSIZE) {
        if ((pte = walk_lookup(pagetable, a)) == 0)
            panic("uvmunmap: walk");
        if ((*pte & PTE_V) == 0)
            panic("uvmunmap: not mapped");
        // 检查是否是叶子页表项（有R/W/X标志位）
        // 如果没有R/W/X标志，说明是指向下一级页表的中间节点，不是叶子
        if ((*pte & (PTE_R | PTE_W | PTE_X)) == 0)
            panic("uvmunmap: not a leaf");
        if (do_free) {
            uint64 pa = PTE2PA(*pte);
            free_page((void *)pa);
        }
        *pte = 0;
    }
}

void
uvmclear(pagetable_t pagetable, uint64 va)
{
    pte_t *pte = walk_lookup(pagetable, va);
    if(pte == 0)
        panic("uvmclear: walk");
    *pte &= ~PTE_U;
}

// 释放用户地址空间的所有物理内存和页表
// pagetable: 用户页表
// sz: 用户地址空间大小（字节）
void uvmfree(pagetable_t pagetable, uint64 sz) {
    if (sz > 0) {
        uvmunmap(pagetable, 0, PGROUNDUP(sz) / PGSIZE, 1);
    }
    destroy_pagetable(pagetable);
}

// allocate and map user memory if process is referencing a page
// that was lazily allocated in sys_sbrk().
// returns 0 if va is invalid or already mapped, or if
// out of physical memory, and physical address if successful.
uint64
vmfault(pagetable_t pagetable, uint64 va, int read)
{
    uint64 mem;
    struct proc *p = myproc();
    pte_t *pte;

    if (va >= p->sz)
        return 0;

    va = PGROUNDDOWN(va);

    // 检查是否已经映射
    pte = walk_lookup(pagetable, va);
    if (pte && (*pte & PTE_V)) {
        return 0;  // 已经映射
    }

    mem = (uint64) alloc_page();
    if (mem == 0)
        return 0;

    memset((void *) mem, 0, PGSIZE);

    if (map_page(p->pagetable, va, mem, PTE_W | PTE_U | PTE_R) != 0) {
        free_page((void *)mem);
        return 0;
    }

    return mem;
}

// 释放用户内存，返回新的大小
uint64
uvmdealloc(pagetable_t pagetable, uint64 oldsz, uint64 newsz)
{
  if(newsz >= oldsz)
    return oldsz;

  if(PGROUNDUP(newsz) < PGROUNDUP(oldsz)){
    int npages = (PGROUNDUP(oldsz) - PGROUNDUP(newsz)) / PGSIZE;
    uvmunmap(pagetable, PGROUNDUP(newsz), npages, 1);
  }

  return newsz;
}




